import cv2
import numpy as np

def nothing(*arg):
    pass

#初始化窗口以及进度条
def Trackbar_Init():
    # 1 create windows
    cv2.namedWindow('h_binary')
    cv2.namedWindow('s_binary')
    cv2.namedWindow('v_binary')

# 2 Create Trackbar
    cv2.createTrackbar('hmin', 'h_binary', 6, 179, nothing)
    cv2.createTrackbar('hmax', 'h_binary', 26, 179, nothing)
    cv2.createTrackbar('smin', 's_binary', 110, 255, nothing)
    cv2.createTrackbar('smax', 's_binary', 255, 255, nothing)
    cv2.createTrackbar('vmin', 'v_binary', 140, 255, nothing)
    cv2.createTrackbar('vmax', 'v_binary', 255, 255, nothing)
    #   创建滑动条     滑动条值名称 窗口名称   滑动条值 滑动条阈值 回调函数


# 在HSV色彩空间下得到二值图
def Get_HSV(image,hmin,hmax,smin,smax,vmin,vmax,color):
    # 1 get trackbar's value


    # 2 to HSV
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
#    cv2.imshow('hsv', hsv)
    h, s, v = cv2.split(hsv)

    # 3 set threshold (binary image)
    # if value in (min, max):white; otherwise:black
    h_binary = cv2.inRange(np.array(h), np.array(hmin), np.array(hmax))
    s_binary = cv2.inRange(np.array(s), np.array(smin), np.array(smax))
    v_binary = cv2.inRange(np.array(v), np.array(vmin), np.array(vmax))

    if(color == "red" or color == "red_exchange"):
        h_binary_2 = cv2.inRange(np.array(h), np.array(hmin), np.array(hmax))
        s_binary_2 = cv2.inRange(np.array(s), np.array(smin), np.array(smax))
        v_binary_2 = cv2.inRange(np.array(v), np.array(vmin), np.array(vmax))



    # 4 get binary（对H、S、V三个通道分别与操作）
    binary = cv2.bitwise_and(h_binary, cv2.bitwise_and(s_binary, v_binary))
    #因为红色有两个范围
    if(color == "red" or color == "red_exchange"):
        binary_2 = cv2.bitwise_and(h_binary_2, cv2.bitwise_and(s_binary_2, v_binary_2))
        binary = cv2.bitwise_and(binary,binary_2)


    # 5 Show

#    cv2.imshow('h_binary', h_binary)
#    cv2.imshow('s_binary', s_binary)
#    cv2.imshow('v_binary', v_binary)
#    cv2.imshow('binary', binary)

    return binary


def Image_Processing_Hough(binary,color):

    # 3 Gausi blur
    blur = cv2.GaussianBlur(binary, (13, 13), 0)

    # 4 Open
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 9))
    Open = cv2.morphologyEx(blur, cv2.MORPH_OPEN, kernel)

    # 5 Close
    Close = cv2.morphologyEx(Open, cv2.MORPH_CLOSE, kernel)

    cv2.imshow(color, Close)
    # 6 Hough Circle detect
    circles = cv2.HoughCircles(Close, cv2.HOUGH_GRADIENT, 1.2, 20, param1=200, param2=90, minRadius=1, maxRadius=300)
    #                                                                     param2:决定圆能否被检测到（越少越容易检测到圆，但相应的也更容易出错）
    # judge if circles is exist


    return circles

def Image_Processing(binary,color):

    # 3 Gausi blur
    blur = cv2.GaussianBlur(binary, (13, 13), 0)

    # 4 Open
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 9))
    Open = cv2.morphologyEx(blur, cv2.MORPH_OPEN, kernel)

    # 5 Close
    Close = cv2.morphologyEx(Open, cv2.MORPH_CLOSE, kernel)

    cv2.imshow(color, Close)

    return Close



def draw_circle(img,circles,b,g,r):

    if type(circles) != type(None):     #防止没有找到圆形轮廓
        
        circles = np.uint16(np.around(circles))

        for i in circles[0,:]:
#            x, y, r = i[0], i[1], i[2]
#            print(x,y,r)
            
            cv2.circle(img, (i[0],i[1]),i[2],(b,g,r),2)  #draw the circle
            cv2.circle(img, (i[0],i[1]),2,(b,g,r),3)


def get_point(circles,x,y,r):

    if type(circles) != type(None):     #防止没有找到圆形轮廓
        
        circles = np.uint16(np.around(circles))

        for i in circles[0,:]:
            if(i[2]>r):
                x=i[0]
                y=i[1] 
                r=i[2]   
        return x,y,r
    else:
        return x,y,r
    

def image_process(binary,color):
        # 3 Gausi blur
    blur = cv2.GaussianBlur(binary, (13, 13), 0)

    # 4 Open
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 9))
    Open = cv2.morphologyEx(blur, cv2.MORPH_OPEN, kernel)

    # 5 Close
    Close = cv2.morphologyEx(Open, cv2.MORPH_CLOSE, kernel)

    edged = cv2.Canny(Close,100,200)

    cv2.imshow(color, edged)
    # 6 Hough Circle detect

    return edged

def findcontours(edged):
    x=0
    y=0
    w=0
    h=0
    contours, hierarchy = cv2.findContours(edged, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    #只找出最大的轮廓
    cnts = sorted(contours,key=cv2.contourArea,reverse = True)

    for c in cnts:
        area = cv2.contourArea(c)
        peri = cv2.arcLength(c,True)
        approx = cv2.approxPolyDP(c,0.02*peri,True)
        if area>20:
            screenCnt = approx

            cv2.drawContours(image,[screenCnt],-1,(0,0,255),2)       
            
            x,y,w,h = cv2.boundingRect(screenCnt)

    return x+w/2,y+h/2
    



def ball_in_box(image):

    gray = cv2.cvtColor(image,cv2.COLOR_BGR2GRAY)
    # 3 Gausi blur
    blur = cv2.GaussianBlur(gray, (13, 13), 0)
    # 4 Open
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 9))
    Open = cv2.morphologyEx(blur, cv2.MORPH_OPEN, kernel)
    # 5 Close
    Close = cv2.morphologyEx(Open, cv2.MORPH_CLOSE, kernel)
    # 6 Hough Circle detect
    circles = cv2.HoughCircles(Close, cv2.HOUGH_GRADIENT, 1.2, 20, param1=200, param2=50, minRadius=1, maxRadius=20)
    #                                                                     param2:决定圆能否被检测到（越少越容易检测到圆，但相应的也更容易出错）


    return circles




#全局变量
#判断框内是否有球
ball_exist = 0
#当exchange_flag为0取红球，为1取蓝球
exchange_flag = 1
x_ball = 0 
y_ball = 0
r_ball = 0

capture = cv2.VideoCapture(1)  # 打开笔记本内置摄像头
while (capture.isOpened()):  # 笔记本内置摄像头被打开后
    retval, image = capture.read(1)  # 从摄像头中实时读取视频
    #img用做绘图，image用作处理
    img = image.copy()

    #找黑球轮廓
    binary_black = Get_HSV(image,0,73,0,255,0,63,"black")

    circle_black = Image_Processing_Hough(binary_black,"black")

    draw_circle(img,circle_black,255,255,255)

    x_ball , y_ball ,r_ball = get_point(circle_black,x_ball,y_ball,r_ball)


    #找黄球轮廓
    binary_yellow = Get_HSV(image,20,27,83,255,155,255,"yellow")

    circle_yellow = Image_Processing_Hough(binary_yellow,"yellow")

    draw_circle(img,circle_yellow,0,255,255)

    x_ball , y_ball ,r_ball = get_point(circle_yellow,x_ball,y_ball,r_ball)



    if(exchange_flag == 0):

        #找红球轮廓
        binary_red = Get_HSV(image,140,179,43,255,16,255,"red")

        circle_red = Image_Processing_Hough(binary_red,"red")

        draw_circle(img,circle_red,0,0,255)

        x_ball , y_ball ,r_ball = get_point(circle_red,x_ball,y_ball,r_ball)

        #找兑换区域
        binary_red_exchange = Get_HSV(image,140,179,43,255,16,255,"red_exchange")
        circle_red_exchange = Image_Processing(binary_red_exchange,"red_exchange")


        x_r_exchange ,y_r_exchange =findcontours(circle_red_exchange)

    else:

        #找蓝球轮廓
        binary_blue = Get_HSV(image,90,115,136,255,20,255,"blue")

        circle_blue = Image_Processing_Hough(binary_blue,"blue")

        draw_circle(img,circle_blue,255,0,0)

        x_ball , y_ball ,r_ball = get_point(circle_blue,x_ball,y_ball,r_ball)


        #找兑换区域
        binary_blue_exchange = Get_HSV(image,90,115,136,255,20,255,"blue_exchange")
        circle_blue_exchange = Image_Processing(binary_blue_exchange,"blue_exchange")

        x_b_exchange ,y_b_exchange =findcontours(circle_blue_exchange)





    #判断球有没有在框中
    box = image[160:320,200:400]

    cv2.imshow("box", box)

    ball_circle = ball_in_box(box)

    draw_circle(box , ball_circle,2,2,2)
    ball_exist = 0

    if type(ball_circle) != type(None):     #防止没有找到圆形轮廓
        ball_exist = 1
    else:
        ball_exist = 0



    cv2.imshow("Video", img)  # 在窗口中显示读取到的视频
    
    key = cv2.waitKey(1)  # 窗口的图像刷新时间为1毫秒
    if key == 27:  # 如果按下esc
        break
capture.release()  # 关闭笔记本内置摄像头
cv2.destroyAllWindows()  # 销毁显示摄像头视频的窗口
